import scipy.io as scio
import numpy as np
import torch
from matplotlib import pyplot as plt


def preprocess(file_path: str = 'data/faces.mat') -> np.ndarray:
    mat_data = scio.loadmat(file_path)
    mat = mat_data['X']

    print('data.shape:', mat.shape)
    return mat


def show_with_grid(filename: str, images: np.ndarray, nrows: int = 10, ncols: int = 10, cmap='grey'):
    """取出 images 中前 nrows * ncols 张图并以网格形式展示在一幅图中，图片将保存到路径 filename。

    Args:
        filename (str): 生成图片的保存路径
        images (np.ndarray): 图像张量
        nrows (int, optional): 网格行数。默认为 10。
        ncols (int, optional): 网格列数。默认为 10。
    """
    fig, axs = plt.subplots(nrows=nrows, ncols=ncols, figsize=(15, 15))
    for i in range(nrows):
        for j in range(ncols):
            axs[i, j].imshow(images[i * nrows + j], cmap=cmap)
            axs[i, j].set_xticks([])
            axs[i, j].set_yticks([])
    plt.savefig(filename, dpi=300, bbox_inches='tight')
    plt.close(fig)


def show_clusters(filename: str, points: torch.Tensor, labels: torch.Tensor, n_cluster: int) -> None:
    """为聚类结果绘制散点图并保存散点图

    Args:
        filename (str): 散点图保存的文件路径。
        points (torch.Tensor): 降维得到的二维点列表（张量），形状应当为 (num_points, 2)。
        labels (torch.Tensor): 与点列表中每个点一一对应的类标签列表，形状应当为 (num_points)。
        n_cluster (int): 聚类数。
    """
    assert points.shape[1] == 2, "Points tensor should have shape (num_points, 2)"
    assert len(points) == len(labels), "Points and labels should have the same length"

    points_ = points.cpu().numpy()
    labels_ = labels.cpu().numpy()

    colors = plt.cm.get_cmap('tab10', n_cluster)

    plt.figure(figsize=(10, 8))
    for i in range(n_cluster):
        cluster_points = points_[labels_ == i]
        plt.scatter(cluster_points[:, 0], cluster_points[:, 1], color=colors(i), label=f'Cluster {i+1}')

    plt.legend()
    plt.savefig(filename)